import csv
import multiprocessing

import torch
from eval import show_without_plt
from eval import check_orthogonality_and_determinant

from optimizer import SGDG
from ffn_getter.get_each_proj import save_proj

GD_W = 1e-2
cayley_alpha = 1e-2
loss_factor = 1e-2

csv_name = "1-8-1"

# 选择计算设备：GPU或CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化随机正交矩阵（旋转矩阵）
def random_orthogonal_matrix(dim, device):
    """生成一个随机的正交矩阵，且行列式值为 1"""
    matrix = torch.randn(dim, dim, device=device)  # 随机生成一个矩阵
    Q, _ = torch.linalg.qr(matrix)  # 使用QR分解生成正交矩阵
    # 修正行列式为 1
    if torch.det(Q) < 0:
        Q[:, 0] = -Q[:, 0]
    return Q


class rotateMatrix(torch.nn.Module):
    def __init__(self, input_dim, device):
        super(rotateMatrix, self).__init__()
        self.input_dim = input_dim
        # self.R = random_orthogonal_matrix(input_dim, device).requires_grad_()  # 确保R可以计算梯度

        self.R = torch.nn.Parameter(random_orthogonal_matrix(input_dim, device))

    def forward(self, X):
        # 旋转输入矩阵 X
        rotated_X = torch.matmul(X, self.R)
        return rotated_X


def compute_loss_l2(R, X, l2_factor=loss_factor):
    """计算 L2 范数损失（基于 R * X）"""
    rotated_X = torch.matmul(X, R)  # 与 R 进行矩阵乘法

    l2_loss = torch.sum(rotated_X ** 2)  # L2 范数损失（平方和）
    total_loss = l2_factor * l2_loss  # 总损失为 L2 范数损失

    return total_loss
def compute_loss_l1(R, X, l1_factor=loss_factor):
    """计算 L1 范数损失（基于 R * X）"""
    rotated_X = torch.matmul(X, R)  # 与R进行矩阵乘法

    l1_loss = torch.sum(torch.abs(rotated_X))  # L1正则化项
    total_loss = l1_factor * l1_loss  # 总损失为 L1 范数损失
    return total_loss

dataset_start = 0
dataset_end = 1000
range_limit = 32
log_step = 1

def load_proj(dataset_start, dataset_end):
    # save_proj(dataset_start, dataset_end, only_target=True)
    arr = []
    for k in range(dataset_start, dataset_end):
        it = torch.load(f'ffn/second_line_input/{k}tensor{0}.pth', weights_only=True).to(torch.float32).to(device)
        arr.append(it)
    return arr

# process = multiprocessing.Process(target=run_save_proj_in_process)
# process.start()
# process.join()

load_fms = load_proj(dataset_start, dataset_end)

rm_dim = load_fms[0].shape[2]

model = rotateMatrix(rm_dim, device).to(device)

optimizer = SGDG(model.parameters(), lr=GD_W,stiefel=True)

fieldnames = ['epoch','dataset_index'] + list( show_without_plt(load_fms[0], print_info=False).keys() ) + ["orthogonal_error", "det"]

def csv_creater(file):
    global fieldnames
    writer = csv.DictWriter(file, fieldnames=fieldnames)
    writer.writeheader()
    file.flush()
    return writer

saved = False

def save_log(epoch, input, orthogonal_error, file,writer, model=model,multi=True):
    global fieldnames
    global saved
    print("------------------------------------")
    print(f"epoch: {epoch}")

    results = {}

    acc_lower_ratio_1e_3 = 0
    acc_lower_ratio_1e_4 = 0

    # result = show_without_plt(input, print_info=True)
    for i in range(len(input)):
        # print('------------')
        if multi:
            result = show_without_plt(input[i] @ model.R, print_info=False)
        else:
            result = show_without_plt(input[i], print_info=False)
        # print('------------')
        result["epoch"] = epoch
        result["orthogonal_error"] = orthogonal_error[0]
        result["det"] = orthogonal_error[1]
        acc_lower_ratio_1e_3 += result["lower_ratio_1e_3"]
        acc_lower_ratio_1e_4 += result["lower_ratio_1e_4"]
        result["dataset_index"] = i



    # print(f"orthogonal_error: {orthogonal_error[0]}")
    # print(f"det: {orthogonal_error[1]}")

    acc_lower_ratio_1e_3 /= len(input)
    acc_lower_ratio_1e_4 /= len(input)
    print(f"acc_lower_ratio_1e_3: {acc_lower_ratio_1e_3}")
    print(f"acc_lower_ratio_1e_4: {acc_lower_ratio_1e_4}")

    results["lower_ratio_1e_3"] = acc_lower_ratio_1e_3
    results["lower_ratio_1e_4"] = acc_lower_ratio_1e_4

    writer.writerow(results)

    # save_csv = file.write(','.join([str(i) for i in arr]) + '\n')
    torch.save(model.R, f'{csv_name}.pth')

    file.flush()
    print("------------------------------------")



print("start training")

file = open(f'{csv_name}.csv', mode='w', newline='')

writer_out = csv_creater(file)

with torch.no_grad():
    save_log(-1, load_fms, check_orthogonality_and_determinant(model.R), file,writer_out ,multi=False)
for epoch in range(2000):
    print(f"epoch: {epoch}")
    for j in range(dataset_start, dataset_end):
        print(f"j: {j}",end=' ')
        input_data = load_fms[j]
        optimizer.zero_grad()

        output = model(input_data)

        total_loss= compute_loss_l1(model.R, input_data)

        total_loss.backward()

        optimizer.step()
    print()
    with torch.no_grad():
        orthogonal_error = check_orthogonality_and_determinant(model.R)

        # 使用 show_without_plt 显示旋转后的矩阵（非plt版本）
    if epoch % log_step == 0:
        with torch.no_grad():
            save_log(epoch, load_fms, orthogonal_error, file,writer_out,model)

    # if epoch % 200 == 0 and epoch != 0:
    #     load_fms = load_proj()

file.close()



